from io import BytesIO

from selenium.webdriver.common.by import By
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.support.ui import WebDriverWait

from PIL import Image


def get_screen_shot(browser):
    screenshot = browser.get_screenshot_as_png()
    image = Image.open(BytesIO(screenshot))
    return image


def get_captcha_pos(browser, css):
    wait = WebDriverWait(browser, 5)
    # 获取验证码元素
    captcha = wait.until(EC.presence_of_element_located((By.CSS_SELECTOR, css)))
    # 获取坐标
    left_x = captcha.location['x']
    left_y = captcha.location['y']

    # 获取大小
    width = captcha.size['width']
    height = captcha.size['height']

    print('pic:width=%d, height=%d' % (width, height))

    right_x = left_x + width
    right_y = left_y + height
    # 比例缩放
    left_x *= 1.25
    left_y *= 1.25
    right_x *= 1.25
    right_y *= 1.25

    return left_x, left_y, right_x, right_y


def crop_img(browser, css, img_name):
    # 获取全图
    full_screen_image = get_screen_shot(browser)
    # 验证图坐标
    left_x, left_y, right_x, right_y = get_captcha_pos(browser, css)
    # 截图并保存
    captcha_img = full_screen_image.crop((left_x, left_y, right_x, right_y))
    captcha_img.save(img_name)
    return captcha_img
